package cn.jbolt.ai.core.chain;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.entity.AiApp;
import cn.jbolt.ai.app.entity.AiAppMessage;
import cn.jbolt.ai.app.entity.AiAppNode;
import cn.jbolt.ai.app.entity.AiAppSession;
import cn.jbolt.ai.app.service.AiAppConfigService;
import cn.jbolt.ai.app.service.AiAppMessageService;
import cn.jbolt.ai.app.service.AiAppService;
import cn.jbolt.ai.app.service.AiAppSessionService;
import cn.jbolt.ai.app.vo.AiChatData;
import cn.jbolt.ai.app.vo.SocketData;
import cn.jbolt.ai.core.chain.node.NodeProviderCenter;
import cn.jbolt.ai.core.chain.node.provider.*;
import cn.jbolt.util.JSONUtil;
import cn.jbolt.util.Result;
import cn.jbolt.util.SpringContextUtil;
//import com.alibaba.fastjson.JSON;
//import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.jboltai.capability.model.JBoltAI;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.chain.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.*;

/**
 * ai app节点链
 */
public class AiAppChain {
    private static final Logger logger = LoggerFactory.getLogger(AiAppChain.class);
    /**
     * 记录每个节点的后置节点,key的格式为 nodeid_case分支名
     */
    private HashMap<Long ,JsonNode> NODE_AFTERS_HOLDER = new HashMap<>();

    private HashMap<Long, String> NODE_TYPE_HOLDER = new HashMap<>();

    private HashMap<Long, JsonNode> NODE_CONFIG_HOLDER = new HashMap<>();

    /**
     * 是否是调试模式
     */
    private boolean isDebug = false;

    /**
     * 开始节点
     */
    private Long[] startIds = new Long[0];

    public AiAppChain(List<AiAppNode> nodes) {

        HashMap<Long, Boolean> hasPrev = new HashMap<>(); //记录每个节点，是否有前置节点
        for (AiAppNode node : nodes) {
            if (hasPrev.containsKey(node.getId()) == false) {
                hasPrev.put(node.getId(), false);
            }
            if (node.getConfig() != null) {
                NODE_CONFIG_HOLDER.put(node.getId(), node.getConfig());
            }
            NODE_TYPE_HOLDER.put(node.getId(), node.getType());
            analyseAfters(node, hasPrev);
        }

        //找没有前置节点的节点，它们就是起始节点
        startIds = hasPrev.entrySet().stream().filter(e -> e.getValue() == false).map(e -> e.getKey()).toArray(Long[]::new);

    }

    /**
     * 启动链条
     * @param data
     * @param session
     */
    public void start(SocketData data, WebSocketSession session) {

        try {
            List<AiAppNode> startNodes = getStartNodes();
            //初始化全局参数
            ChainContext context = initContext(data, session);


            Chain chain = JBoltAI.chain();
            chain.setId(data.getFlag());
            HashMap<Long, ChainNode> nodeCache = new HashMap<>();
            for (AiAppNode node : startNodes) {
                ChainNode chainNode = prepareNode(node, context, nodeCache);
                chain.addNodes(chainNode);
            }
            chain.onFail((e, error) -> {
                logger.error("事件编排执行失败：{}", error.getFullMsg());

            });
            chain.onComplete(e -> {
                context.sendOverMsg();

                switch (e.getState()) {
                    case SUCCESS:
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_SUCCESS);
                        break;
                    case CANCEL:
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_CANCEL);
                        break;
                    case FAIL:
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_ERROR);
                        context.getResponseMsg().setErrorMsg(e.getError() != null ? e.getError().getMsg() : "AI回答失败...");
                        break;
                }
                SpringContextUtil.getBean(AiAppMessageService.class).update(context.getResponseMsg());

            });

            chain.publish();
        } catch (Exception e) {
            e.printStackTrace();
            SocketData res = SocketData.copy(data);
            res.setAction(SocketData.ACTION_ERROR);
            res.setData(JSONUtil.createObjectNode().put("msg", "程序异常"));
            try {
                session.sendMessage(new TextMessage(res.toString()));
                res.setAction(SocketData.ACTION_OVER).setData(null);
                session.sendMessage(new TextMessage(res.toString()));

            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        }


    }

    /**
     * 准备节点
     * @param node
     * @param nodeCache
     * @return
     */
    private ChainNode prepareNode(AiAppNode node, ChainContext context, HashMap<Long, ChainNode> nodeCache) {

        ChainNode chainNode = nodeCache.get(node.getId());
        if (chainNode != null) {
            return chainNode;
        }


        Object nodeProvider = NodeProviderCenter.getNodeProvider(node.getType());
        if (nodeProvider == null) {
            throw new RuntimeException("未注册的节点类型: " + node.getType());
        }
        AiAppConfigService configService = SpringContextUtil.getBean(AiAppConfigService.class);
        JsonNode config = NODE_CONFIG_HOLDER.containsKey(node.getId()) ? NODE_CONFIG_HOLDER.get(node.getId()) : configService.getNodeValidConfig(context.getAiApp().getId(), context.getAiSession().getId(), node.getId());

        if (nodeProvider instanceof EventNodeProvider) {
            //事件节点
            chainNode = new EventChainNode(c -> {
                try {

                    return ((EventNodeProvider)nodeProvider).provide(config, context);
                } catch (Exception e) {
                    c.fail(EventErrorType.PROGRAM_EXCEPTION, e);
                    return null;
                }
            });
            nodeCache.put(node.getId(), chainNode);
            for (AiAppNode nextNode : getNextNodes(node.getId())) {
                prepareNode(nextNode, context, nodeCache).addPrevNode(chainNode);
            }
            for (AiAppNode nextNode : getNextNodes(node.getId(), ChainNode.AFTER_KEY_FAIL)) {
                prepareNode(nextNode, context, nodeCache).addPrevNodeWhenFail(chainNode);
            }
        } else if (nodeProvider instanceof CommonNodeProvider) {
            //通用节点
            chainNode = new CommonChainNode(c -> {
                ((CommonNodeProvider)nodeProvider).provide(config, context);
            });
            nodeCache.put(node.getId(), chainNode);
            for (AiAppNode nextNode : getNextNodes(node.getId())) {
                prepareNode(nextNode, context, nodeCache).addPrevNode(chainNode);
            }
            for (AiAppNode nextNode : getNextNodes(node.getId(), ChainNode.AFTER_KEY_FAIL)) {
                prepareNode(nextNode, context, nodeCache).addPrevNodeWhenFail(chainNode);
            }
        } else if (nodeProvider instanceof CaseNodeProvider) {
            //通用节点
            chainNode = new CaseChainNode(c -> {
                return ((CaseNodeProvider)nodeProvider).provide(config, context);
            });
            nodeCache.put(node.getId(), chainNode);
            HashMap<String, List<AiAppNode>> allNextNodes = getAllNextNodes(node.getId());
            for (Map.Entry<String, List<AiAppNode>> entry : allNextNodes.entrySet()) {
                for (AiAppNode nextNode : entry.getValue()) {
                    prepareNode(nextNode, context, nodeCache).addPrevWhenCase((CaseChainNode)chainNode, entry.getKey() );
                }
            }
        } else if (nodeProvider instanceof BooleanNodeProvider) {
            //布尔节点
            chainNode = new BooleanChainNode(c -> {
                return ((BooleanNodeProvider)nodeProvider).provide(config, context);
            });
            nodeCache.put(node.getId(), chainNode);
            for (AiAppNode nextNode : getNextNodes(node.getId(), BooleanChainNode.AFTER_KEY_TRUE)) {
                prepareNode(nextNode, context, nodeCache).addPrevWhenTrue((BooleanChainNode)chainNode);
            }
            for (AiAppNode nextNode : getNextNodes(node.getId(), BooleanChainNode.AFTER_KEY_FALSE)) {
                prepareNode(nextNode, context, nodeCache).addPrevWhenFalse((BooleanChainNode)chainNode);
            }
            for (AiAppNode nextNode : getNextNodes(node.getId(), ChainNode.AFTER_KEY_FAIL)) {
                prepareNode(nextNode, context, nodeCache).addPrevNodeWhenFail(chainNode);
            }
        } else if (nodeProvider instanceof IntentionNodeProvider) {
            //意图识别节点
            chainNode = new IntentionChainNode(c -> {
                return ((IntentionNodeProvider)nodeProvider).provide(config, context);
            });
            nodeCache.put(node.getId(), chainNode);

            HashMap<String, List<AiAppNode>> allNextNodes = getAllNextNodes(node.getId());
            for (Map.Entry<String, List<AiAppNode>> entry : allNextNodes.entrySet()) {
                for (AiAppNode nextNode : entry.getValue()) {
                    prepareNode(nextNode, context, nodeCache).addPrevWhenIntentionMatch(entry.getKey(), (IntentionChainNode)chainNode);
                }
            }

        }


        return chainNode;
    }

    /**
     * 初始化上下文
     * @param data
     * @param session
     * @return
     */
    private ChainContext initContext(SocketData data, WebSocketSession session) {
        Result appResult = SpringContextUtil.getBean(AiAppService.class).getAppById(data.getApp());
        Result sessionResult = SpringContextUtil.getBean(AiAppSessionService.class).getSession(data.getSessionId());
        AiAppSession aiSession = (AiAppSession)sessionResult.getData();
        AiApp aiApp = (AiApp) appResult.getData();

        ChainContext context = new ChainContext();
        context.setSocketData(data);
        context.setSocketSession(session);

        context.setAiApp(aiApp);
        context.setAiSession(aiSession);

        switch (data.getAction()) {
            case SocketData.ACTION_REQUEST: {

                initRequestData(context);
                break;
            }
            case SocketData.ACTION_REGENERATE:
                initRegenerateData(context);
                break;
        }



        return context;
    }

    /**
     * 初始化重新生成需要的数据
     * @param context
     */
    private static void initRegenerateData(ChainContext context) {
        Result<AiAppMessage[]> result = SpringContextUtil.getBean(AiAppMessageService.class)
                .regenerateMsg(context.getSocketData().getMsgId(),
                        "思考中...");
        if (result.isError()) {
            context.sendErrorMsg("程序发生错误：" + result.getMsg());
            throw new RuntimeException(result.getMsg() );
        }
        context.setRequestMsg(result.getData()[0]);
        context.setResponseMsg(result.getData()[1]);
        context.sendSyncDataMsg();
        AiChatData data = new AiChatData();
        data.setRole("user");
        data.setContent(new AiChatData.AiChatContent(context.getRequestMsg().getContent()));
        context.getSocketData().setData(data);
    }


    private static void initRequestData(ChainContext context) {
        //生成对应的提问跟回答的消息
        Result<AiAppMessage[]> result = SpringContextUtil.getBean(AiAppMessageService.class)
                .addMessageReturnAnswer(context.getSocketData(),
                AiAppMessage.DEFAULT_MSG, context.getAiApp(), context.getAiSession().getId());
        if (result.isError()) {
            context.sendErrorMsg("程序发生错误：" + result.getMsg());
            throw new RuntimeException(result.getMsg() );
        }
        context.setRequestMsg(result.getData()[0]);
        context.setResponseMsg(result.getData()[1]);
        context.sendSyncDataMsg();
    }


    /**
     * 解析后置节点，将他们缓存起来
     * @param node
     */
    private void analyseAfters(AiAppNode node, HashMap<Long, Boolean> hasPrev) {
        String afters = node.getAfters();
        if (StrUtil.isBlank(afters)) return;
        Long[] afterIds = null;
        if (JSONUtil.isTypeJSONArray(afters)) {
            //格式为 ["1", "2"]
            List<Long> afterIdList = JSONUtil.toObjectList(afters, Long.class);
            afterIds = afterIdList.toArray(new Long[0]);
            ObjectNode objectNode = JSONUtil.createObjectNode();
            objectNode.set(ChainNode.AFTER_KEY_SUCCESS, JSONUtil.toJsonArray(afters));
            NODE_AFTERS_HOLDER.put(node.getId(), objectNode);
            for (Long afterId : afterIds) {
                hasPrev.put(afterId, true);
            }
        } else {
            // 格式为 {“true":["1", "2"], "false":["3","4"]}
            ObjectNode aftersObj = JSONUtil.toJsonObject(afters);
            NODE_AFTERS_HOLDER.put(node.getId(), aftersObj);
            Iterator<String> keys = aftersObj.fieldNames();
            while (keys.hasNext()) {
                String key = keys.next(); // 获取当前键名
                JsonNode jsonNode = aftersObj.get(key);
                List<Long> afterIdList = JSONUtil.toObjectList(JSONUtil.toJson(jsonNode), Long.class);

                afterIds = afterIdList.toArray(new Long[0]);
                for (Long afterId : afterIds) {
                    hasPrev.put(afterId, true);
                }
            }
        }

    }


    /**
     * 获取起始节点
     * @return
     */
    public List<AiAppNode> getStartNodes() {
        List<AiAppNode> nodes = new ArrayList<>();
        for (Long startId : startIds) {
            AiAppNode node = new AiAppNode();
            node.setId(startId);
            node.setType(NODE_TYPE_HOLDER.get(startId));
            nodes.add(node);
        }
        return nodes;
    }

    /**
     * 获取后续节点，适合有分支的节点
     * @param fromNode
     * @param caseName
     * @return
     */
    public List<AiAppNode> getNextNodes(Long fromNode, String caseName) {
        JsonNode afters = NODE_AFTERS_HOLDER.get(fromNode);

        if (afters == null || afters.has(caseName) == false) {
            return Collections.emptyList();
        }

        JsonNode jsonNode = afters.get(caseName);
        List<Long> javaList = JSONUtil.toObjectList(JSONUtil.toJson(jsonNode), Long.class);

        Long[] afterIds = javaList.toArray(new Long[0]);
        if (afterIds == null) {
            return new ArrayList<>();
        }
        List<AiAppNode> nodes = new ArrayList<>();
        for (Long afterId : afterIds) {
            AiAppNode node = new AiAppNode();
            node.setId(afterId);
            node.setType(NODE_TYPE_HOLDER.get(afterId));
            nodes.add(node);
        }
        return nodes;
    }


    /**
     * 获取所有后续节点
     * @param fromNode
     * @return
     */
    public HashMap<String, List<AiAppNode>> getAllNextNodes(Long fromNode) {
        HashMap<String, List<AiAppNode>> map = new HashMap<>();
        JsonNode afters = NODE_AFTERS_HOLDER.get(fromNode);
        if (afters == null) {
            return map;
        }

        Iterator<String> keys = afters.fieldNames();
        while (keys.hasNext()) {
            String key = keys.next(); // 获取当前键名
            JsonNode jsonNode = afters.get(key);
            List<Long> afterIds = JSONUtil.toObjectList(JSONUtil.toJson(jsonNode), Long.class);
            if (afterIds == null) {
                continue;
            }
            List<AiAppNode> nodes = new ArrayList<>();
            for (Long afterId : afterIds) {
                AiAppNode node = new AiAppNode();
                node.setId(afterId);
                node.setType(NODE_TYPE_HOLDER.get(afterId));
                nodes.add(node);
            }
            map.put(key, nodes);
        }
        return map;
    }


    /**
     * 获取后续节点，适合无分支的节点
     * @param fromNode
     * @return
     */
    public List<AiAppNode> getNextNodes(Long fromNode) {
        return getNextNodes(fromNode, ChainNode.AFTER_KEY_SUCCESS);
    }


    public boolean isDebug() {
        return isDebug;
    }

    public void setDebug(boolean debug) {
        isDebug = debug;
    }
}
